package cn.smartjavaai.action.model;

import ai.djl.modality.Classifications;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.util.NDImageUtils;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDList;
import ai.djl.ndarray.NDManager;
import ai.djl.ndarray.types.DataType;
import ai.djl.ndarray.types.Shape;
import ai.djl.translate.Batchifier;
import ai.djl.translate.Translator;
import ai.djl.translate.TranslatorContext;
import cn.smartjavaai.common.utils.LetterBoxUtils;

import java.util.Arrays;

/**
 * @author dwj
 */
public class CommonActionTranslator implements Translator<Image, Classifications> {

    private int width = 224;

    private int height = 224;

    private String[] labels;

    public static String rawLabels = "'abseiling','air_drumming','answering_questions','applauding','applying_cream','archery'," +
            "'arm_wrestling','arranging_flowers','assembling_computer','auctioning','baby_waking_up','baking_cookies'," +
            "'balloon_blowing','bandaging','barbequing','bartending','beatboxing','bee_keeping','belly_dancing'," +
            "'bench_pressing','bending_back','bending_metal','biking_through_snow','blasting_sand','blowing_glass'," +
            "'blowing_leaves','blowing_nose','blowing_out_candles','bobsledding','bookbinding','bouncing_on_trampoline'," +
            "'bowling','braiding_hair','breading_or_breadcrumbing','breakdancing','brush_painting','brushing_hair'," +
            "'brushing_teeth','building_cabinet','building_shed','bungee_jumping','busking','canoeing_or_kayaking'," +
            "'capoeira','carrying_baby','cartwheeling','carving_pumpkin','catching_fish','catching_or_throwing_baseball'," +
            "'catching_or_throwing_frisbee','catching_or_throwing_softball','celebrating','changing_oil','changing_wheel'," +
            "'checking_tires','cheerleading','chopping_wood','clapping','clay_pottery_making','clean_and_jerk'," +
            "'cleaning_floor','cleaning_gutters','cleaning_pool','cleaning_shoes','cleaning_toilet','cleaning_windows'," +
            "'climbing_a_rope','climbing_ladder','climbing_tree','contact_juggling','cooking_chicken','cooking_egg'," +
            "'cooking_on_campfire','cooking_sausages','counting_money','country_line_dancing','cracking_neck','crawling_baby'," +
            "'crossing_river','crying','curling_hair','cutting_nails','cutting_pineapple','cutting_watermelon'," +
            "'dancing_ballet','dancing_charleston','dancing_gangnam_style','dancing_macarena','deadlifting'," +
            "'decorating_the_christmas_tree','digging','dining','disc_golfing','diving_cliff','dodgeball','doing_aerobics'," +
            "'doing_laundry','doing_nails','drawing','dribbling_basketball','drinking','drinking_beer','drinking_shots'," +
            "'driving_car','driving_tractor','drop_kicking','drumming_fingers','dunking_basketball','dying_hair'," +
            "'eating_burger','eating_cake','eating_carrots','eating_chips','eating_doughnuts','eating_hotdog'," +
            "'eating_ice_cream','eating_spaghetti','eating_watermelon','egg_hunting','exercising_arm'," +
            "'exercising_with_an_exercise_ball','extinguishing_fire','faceplanting','feeding_birds','feeding_fish'," +
            "'feeding_goats','filling_eyebrows','finger_snapping','fixing_hair','flipping_pancake','flying_kite'," +
            "'folding_clothes','folding_napkins','folding_paper','front_raises','frying_vegetables','garbage_collecting'," +
            "'gargling','getting_a_haircut','getting_a_tattoo','giving_or_receiving_award','golf_chipping','golf_driving'," +
            "'golf_putting','grinding_meat','grooming_dog','grooming_horse','gymnastics_tumbling','hammer_throw'," +
            "'headbanging','headbutting','high_jump','high_kick','hitting_baseball','hockey_stop','holding_snake'," +
            "'hopscotch','hoverboarding','hugging','hula_hooping','hurdling','hurling_-sport-','ice_climbing','ice_fishing'," +
            "'ice_skating','ironing','javelin_throw','jetskiing','jogging','juggling_balls','juggling_fire'," +
            "'juggling_soccer_ball','jumping_into_pool','jumpstyle_dancing','kicking_field_goal','kicking_soccer_ball'," +
            "'kissing','kitesurfing','knitting','krumping','laughing','laying_bricks','long_jump','lunge','making_a_cake'," +
            "'making_a_sandwich','making_bed','making_jewelry','making_pizza','making_snowman','making_sushi','making_tea'," +
            "'marching','massaging_back','massaging_feet','massaging_legs','massaging_person's_head','milking_cow'," +
            "'mopping_floor','motorcycling','moving_furniture','mowing_lawn','news_anchoring','opening_bottle'," +
            "'opening_present','paragliding','parasailing','parkour','passing_American_football_-in_game-'," +
            "'passing_American_football_-not_in_game-','peeling_apples','peeling_potatoes','petting_animal_-not_cat-'," +
            "'petting_cat','picking_fruit','planting_trees','plastering','playing_accordion','playing_badminton'," +
            "'playing_bagpipes','playing_basketball','playing_bass_guitar','playing_cards','playing_cello','playing_chess'," +
            "'playing_clarinet','playing_controller','playing_cricket','playing_cymbals','playing_didgeridoo','playing_drums'," +
            "'playing_flute','playing_guitar','playing_harmonica','playing_harp','playing_ice_hockey','playing_keyboard'," +
            "'playing_kickball','playing_monopoly','playing_organ','playing_paintball','playing_piano','playing_poker'," +
            "'playing_recorder','playing_saxophone','playing_squash_or_racquetball','playing_tennis','playing_trombone'," +
            "'playing_trumpet','playing_ukulele','playing_violin','playing_volleyball','playing_xylophone','pole_vault'," +
            "'presenting_weather_forecast','pull_ups','pumping_fist','pumping_gas','punching_bag','punching_person_-boxing-'," +
            "'push_up','pushing_car','pushing_cart','pushing_wheelchair','reading_book','reading_newspaper','recording_music'," +
            "'riding_a_bike','riding_camel','riding_elephant','riding_mechanical_bull','riding_mountain_bike','riding_mule'," +
            "'riding_or_walking_with_horse','riding_scooter','riding_unicycle','ripping_paper','robot_dancing','rock_climbing'," +
            "'rock_scissors_paper','roller_skating','running_on_treadmill','sailing','salsa_dancing','sanding_floor'," +
            "'scrambling_eggs','scuba_diving','setting_table','shaking_hands','shaking_head','sharpening_knives'," +
            "'sharpening_pencil','shaving_head','shaving_legs','shearing_sheep','shining_shoes','shooting_basketball'," +
            "'shooting_goal_-soccer-','shot_put','shoveling_snow','shredding_paper','shuffling_cards','side_kick'," +
            "'sign_language_interpreting','singing','situp','skateboarding','ski_jumping','skiing_-not_slalom_or_crosscountry-'," +
            "'skiing_crosscountry','skiing_slalom','skipping_rope','skydiving','slacklining','slapping','sled_dog_racing'," +
            "'smoking','smoking_hookah','snatch_weight_lifting','sneezing','sniffing','snorkeling','snowboarding','snowkiting'," +
            "'snowmobiling','somersaulting','spinning_poi','spray_painting','spraying','springboard_diving','squat'," +
            "'sticking_tongue_out','stomping_grapes','stretching_arm','stretching_leg','strumming_guitar','surfing_crowd'," +
            "'surfing_water','sweeping_floor','swimming_backstroke','swimming_breast_stroke','swimming_butterfly_stroke'," +
            "'swing_dancing','swinging_legs','swinging_on_something','sword_fighting','tai_chi','taking_a_shower','tango_dancing'," +
            "'tap_dancing','tapping_guitar','tapping_pen','tasting_beer','tasting_food','testifying','texting','throwing_axe'," +
            "'throwing_ball','throwing_discus','tickling','tobogganing','tossing_coin','tossing_salad','training_dog'," +
            "'trapezing','trimming_or_shaving_beard','trimming_trees','triple_jump','tying_bow_tie','tying_knot_-not_on_a_tie-'," +
            "'tying_tie','unboxing','unloading_truck','using_computer','using_remote_controller_-not_gaming-','using_segway'," +
            "'vault','waiting_in_line','walking_the_dog','washing_dishes','washing_feet','washing_hair','washing_hands'," +
            "'water_skiing','water_sliding','watering_plants','waxing_back','waxing_chest','waxing_eyebrows','waxing_legs'," +
            "'weaving_basket','welding','whistling','windsurfing','wrapping_present','wrestling','writing','yawning','yoga','zumba'";

    public CommonActionTranslator(int width, int height) {
        this.width = width;
        this.height = height;
    }

    @Override
    public void prepare(TranslatorContext ctx) throws Exception {
        labels = rawLabels.replace("'", "").split(",");
        labels = Arrays.stream(labels)
                .map(String::trim)
                .toArray(String[]::new);
    }

    @Override
    public NDList processInput(TranslatorContext ctx, Image input) {
        NDManager manager = ctx.getNDManager();
        NDArray array = input.toNDArray(manager, Image.Flag.COLOR); // HWC, uint8
        array = NDImageUtils.resize(array, width, height);
        // 转 float32
        array = array.toType(DataType.FLOAT32, false);
        // 自己归一化，注意均值和标准差要乘以255，因为原始是0~255
        float[] mean = {0.485f * 255, 0.456f * 255, 0.406f * 255};
        float[] std = {0.229f * 255, 0.224f * 255, 0.225f * 255};
        // 增加 batch 维度，变成 (1, H, W, C)
        array = array.expandDims(0);
        return new NDList(array);
    }

    @Override
    public Classifications processOutput(TranslatorContext ctx, NDList list) {
        NDArray output = list.singletonOrThrow();
        output = output.softmax(1); // 计算概率
        Classifications classifications = new Classifications(Arrays.asList(labels), output);
        return classifications;

    }

    @Override
    public Batchifier getBatchifier() {
        return null;
    }




}
